package tool.code
 
trait CodeGenerator {
	self: Reporter =>
	import parser.Trees._
	import analyzer.Symbols._
	import analyzer.Types._
	import cafebabe._

	import scala.collection.mutable.HashMap
	
	import AbstractByteCodes._
	import ByteCodes._

	var outputDir : String = ""
	var classFile : ClassFile = null
	var codeHandler : CodeHandler = null

	var currentClass: ClassSymbol = null
	var currentMethod: MethodSymbol = null

	var varToIndex: HashMap[String,Int] = new HashMap[String,Int]()

    def isMember(i: Identifier) = {
		currentMethod.lookupVar(i.value) match{
			case Some(p) => false
			case None => currentClass.lookupVar(i.value) match {
			  case Some(q) => true
              case None => false
			}
		}
	}

	def getIndex(s: String) = varToIndex get(s) match {
	  case Some(p) => p
      case None => throw new Exception("This shouldn't happen, if you find the programmer, please whack him")
	}

	def addVar(s: String) = varToIndex += new Pair(s, codeHandler.getFreshVar)

	def freeVar(s: String) = {
		varToIndex get (s) match{
			case Some(p) => codeHandler.freeVar(p); varToIndex -= s
			case None =>
		}
	}

	/** Writes the proper .class file in a given directory. An empty string for dir is equivalent to "./". */
	def generateClassFile(gs: GlobalScope, ct: Class, dir: String): Unit = {
	    
        var varDecls: List[VarDecl] = Nil
        var metDecls: List[MethDecl] = Nil
        var name: String = null

		currentClass = ct.getSymbol
        currentClass.members
        
		ct match {
		    case ClassDecl(id,vDecls,mDecls) =>
		        classFile = new ClassFile(id.value, None)
                varDecls = vDecls;
                metDecls = mDecls
                name = id.value
		    case ClassWithExtends(id, superClass, vDecls, mDecls) =>
		        classFile = new ClassFile(id.value,Some(superClass.value))
                varDecls = vDecls;
                metDecls = mDecls
                name = id.value
		}
  
        // Generate default constructor
        classFile.addDefaultConstructor
  
		varDecls.foreach(g =>
			 classFile.addField(generateTypeCodeSignature(g.id.getSymbol.asInstanceOf[VariableSymbol].getType), g.id.value)
		)

		metDecls.foreach(g =>
			generateMethod(g)
		)
        
        classFile.writeToFile(dir + name + ".class")
        outputDir = dir
	}

	def generateMethod(metDecl: MethDecl) = {
		val met = metDecl.getSymbol
		currentMethod = met;
		var index = 1 // method argument indexes start from 1
		var encodedRetType = ""

		varToIndex.clear()

		met.argList.foreach(g =>
			encodedRetType += generateTypeCodeSignature(g.getType)
		)
		codeHandler = classFile.addMethod(generateTypeCodeSignature(met.getType), met.name, encodedRetType).codeHandler

		met.argList.foreach(g => 
			{varToIndex += new Pair(g.name, index); index += 1}
		)

		met.members.keys.foreach(g => 
		    varToIndex += new Pair(g, codeHandler.getFreshVar)
		)	

		generateStatementCode(metDecl.sts)
  
		generateExpressionCode(metDecl.returnExpr)
  
		met.getType match{
			case TInt | TBool => codeHandler << IRETURN
			case TArray | TString | TObject(_) => codeHandler << ARETURN
			case _ => // shouldn't happen, but makes the compiler happy
		}
 
        varToIndex.keys.foreach(v => freeVar(v));

		codeHandler.freeze
	}

	def generateMainObjectClassFile(main: MainObject, gs: GlobalScope, dir: String): Unit = {
		outputDir = dir
		classFile = new ClassFile(gs.mainClass.name, None)
		classFile.addDefaultConstructor
		codeHandler = classFile.addMainMethod.codeHandler

		generateStatementCode(main.stat)
		codeHandler << RETURN
		codeHandler.freeze

		classFile.writeToFile(outputDir + gs.mainClass.name + ".class")
	}

	def generateStatementCode(stat: StatTree): Unit = stat match{
		case Block(stats) => 
			stats.foreach(g => generateStatementCode(g))
		case While(t, stat) => 
			val end_of_loop = codeHandler.getFreshLabel("end_of_loop")
			val loop = codeHandler.getFreshLabel("loop")
			codeHandler << Label(loop)
			generateExpressionCode(t)
			codeHandler << IfEq(end_of_loop)
			generateStatementCode(stat)
			codeHandler << Goto(loop)
			codeHandler << Label(end_of_loop)
		case If(t, s) => 
			val after = codeHandler.getFreshLabel("after")
			generateExpressionCode(t)
			codeHandler << IfEq(after)
			generateStatementCode(s)
			codeHandler << Label(after)
		case IfElse(t, s1, s2) =>
			val not_true = codeHandler.getFreshLabel("not_true")
			val after = codeHandler.getFreshLabel("after_true")
			generateExpressionCode(t)
			codeHandler << IfEq(not_true)
			generateStatementCode(s1)
			codeHandler << Goto(after)
			codeHandler << Label(not_true)
			generateStatementCode(s2)
			codeHandler << Label(after)
		case Println(e) => 
			codeHandler << GetStatic("java/lang/System", "out", "Ljava/io/PrintStream;")
			generateExpressionCode(e)
			codeHandler << InvokeVirtual("java/io/PrintStream", "println","("+generateTypeCode(e.getType)+")V")
		case Assign(id, e) => {
			generateStore(id,e)}
		case AssignArrayCell(id, idx, v) =>
			generateLoad(id)
			generateExpressionCode(idx)
			generateExpressionCode(v)
			codeHandler << IASTORE

		case TryCatchBlock(b, cs, f) =>
			val tryStart = codeHandler.getFreshLabel("tryStart")
			val tryEnd = codeHandler.getFreshLabel("tryEnd")
			val catchesEnd = codeHandler.getFreshLabel("catchesEnd")
			val normalBlockEnd = codeHandler.getFreshLabel("normalBlockEnd")
			val finallyStart = codeHandler.getFreshLabel("finallyStart")
                 
			// Generate try block
			codeHandler << Label(tryStart)
			generateStatementCode(b)
			codeHandler << Label(tryEnd)
			codeHandler << Goto(normalBlockEnd)
   
			// Generate catch blocks
			cs.foreach(g => {
				val catchStart = codeHandler.getFreshLabel("catchStart")
				varToIndex += new Pair(g.exceptionId.value, codeHandler.getFreshVar)
				codeHandler << Label(catchStart)
				codeHandler << AStore(getIndex(g.exceptionId.value))
				generateStatementCode(g.body)
				freeVar(g.exceptionId.value)
				codeHandler << Goto(normalBlockEnd)
				codeHandler.addHandler(tryStart, tryEnd, catchStart, Some(g.exceptionType.value))
			})
        
			// Generate finally if exists
			f match {
				case None => codeHandler << Label(normalBlockEnd)
				case Some(fStmt) => 
					codeHandler << Label(catchesEnd)
   
					// Generate false "catch" for finally
					val exeptionId = getUniqueId("exeptionToRaise")
					val finallyCatchStart = codeHandler.getFreshLabel("finallyCatchStart")
					codeHandler << Label(finallyCatchStart)
					varToIndex += new Pair(exeptionId, codeHandler.getFreshVar) 
					codeHandler << AStore(getIndex(exeptionId)) // Store thrown exception
					codeHandler << Jsr(finallyStart)
					codeHandler << ALoad(getIndex(exeptionId)) // Put thrown exception back to the stack
					codeHandler << ATHROW // Re-throw the exception
					codeHandler.addHandler(tryStart, catchesEnd, finallyCatchStart,None) // any

					// Generate finally block
					codeHandler << Label(finallyStart)
					generateStatementCode(fStmt)
					freeVar(exeptionId)

					// Jump to finally subroutine
					codeHandler << Label(normalBlockEnd)
					codeHandler << Jsr(finallyStart)
			}
		case Throw(expr) =>
			generateExpressionCode(expr)
			codeHandler << ATHROW
        case Finally(body) =>
            val returnAddressId = getUniqueId("returnAddress")
            varToIndex += new Pair(returnAddressId, codeHandler.getFreshVar)
            codeHandler << AStore(getIndex(returnAddressId))
            generateStatementCode(body)
            codeHandler << Ret(getIndex(returnAddressId))
            freeVar(returnAddressId)
	}
 
    private def getUniqueId(baseId : String) : String = {
      var suffix = 0
      
      while (true) {
        varToIndex.get(baseId+suffix) match {
          case Some(s) => suffix += 1
          case None => return baseId+suffix
        }
      }
      
      return baseId+suffix
    }

	def isSimplePrintableValue(e: ExprTree): Boolean = e match{
		case StringLit(_) => true
		case Identifier(_) => true 
		case MethodCall(_, _, _) => true
		case _ => 
			e.getType match{
				case TBool | TInt => true
				case _ => false
			}
	}

	def generateCorrectAppend(e: ExprTree): Unit = {
		e.getType match {
			case TInt | TBool => 
				codeHandler << InvokeVirtual("java/lang/StringBuilder", "append", "(I)Ljava/lang/StringBuilder;")
			case TString =>
				codeHandler << InvokeVirtual("java/lang/StringBuilder", "append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;")
			case _ => // shouldn't happen, but makes the compiler happy
		}
	}

	var stringBuilderInitted: Boolean = false

	def createStringBuilder(expr: ExprTree) = {
		codeHandler << DefaultNew("java/lang/StringBuilder")
		stringBuilderInitted = true
		generateExpressionCode(expr)
		codeHandler << InvokeVirtual("java/lang/StringBuilder", "toString", "()Ljava/lang/String;")
		stringBuilderInitted = false
	}

	def generateExpressionCode(expr: ExprTree): Unit = expr match{
		case IntLit(i) => codeHandler << Ldc(i)
		case BoolLit(b) => codeHandler << Ldc(if(b) 1 else 0)
		case StringLit(s) => codeHandler << Ldc(s)

		case Plus(lhs, rhs) if expr.getType == TString && stringBuilderInitted =>
			generateExpressionCode(lhs); if(isSimplePrintableValue(lhs)) generateCorrectAppend(lhs)
			generateExpressionCode(rhs); if(isSimplePrintableValue(rhs)) generateCorrectAppend(rhs)
		case Plus(lhs, rhs) if expr.getType == TString =>
			createStringBuilder(expr)
		case Plus(lhs, rhs) =>
			generateExpressionCode(lhs); generateExpressionCode(rhs)
			codeHandler << IADD
		case Minus(lhs, rhs) =>
			generateExpressionCode(lhs); generateExpressionCode(rhs)
			codeHandler << ISUB
		case Times(lhs, rhs) =>
			generateExpressionCode(lhs); generateExpressionCode(rhs)
			codeHandler << IMUL
		case Div(lhs, rhs) =>
			generateExpressionCode(lhs); generateExpressionCode(rhs)
			codeHandler << IDIV
		case And(lhs, rhs) =>
			generateConditionValue(expr)
		case Or(lhs, rhs) =>
			generateConditionValue(expr)
		case Equals(lhs, rhs) =>
			generateConditionValue(expr)
		case LessThan(lhs, rhs) =>
			generateConditionValue(expr)
		case Not(v) =>
			generateConditionValue(expr)
		case ArrayCell(t, idx) =>
			generateExpressionCode(t)
			generateExpressionCode(idx)
			codeHandler << IALOAD
		case Length(tab) =>
			generateExpressionCode(tab)
			codeHandler << ARRAYLENGTH
		case m @ MethodCall(inst, meth, args) =>
			generateExpressionCode(inst)
			args.foreach(g => generateExpressionCode(g))
            codeHandler << InvokeVirtual(generateTypeCode(inst.getType), meth.value, generateMethodSignature(m))
		case ArrayCreation(s) =>
			generateExpressionCode(s)
			codeHandler << NewArray(10) // Array of int
		case Instanciation(t) =>
			codeHandler << DefaultNew(t.value)
		case ident @ Identifier(v) => 
		    if (v.equals("this")) {
		        codeHandler << ALoad(0)
		    } else {
		        generateLoad(ident)
		    }
		case _ =>
	}

	def generateMethodSignature(met: MethodCall) = {
		var t = "("
		met.methodId.getSymbol.asInstanceOf[MethodSymbol].argList.foreach(g =>
		    t += generateTypeCodeSignature(g.getType)
		)
		t += ")"
		t += generateTypeCodeSignature(met.getType)
    
		t
	}

	/**
	 * Generates an int value from a binary test (equals or lessthan) 
	 * The following method may needs to be modified so as to reduce the size of the result
	 */
	def generateConditionValue(test: ExprTree): Unit = {
		val zero_label = codeHandler.getFreshLabel("zero_label")
		val after_label = codeHandler.getFreshLabel("after_label")

		test match{
			case BoolLit(v) => 
				codeHandler << Ldc(if(v) 1 else 0)
				return

			case Equals(lhs, rhs) =>
				generateExpressionCode(lhs); generateExpressionCode(rhs)
				lhs.getType match{
					case TString | TObject(_) | TArray => codeHandler << If_ACmpNe(zero_label)
					case _ => codeHandler << If_ICmpNe(zero_label)
				}
			case LessThan(lhs, rhs) =>
				generateExpressionCode(lhs); generateExpressionCode(rhs)
				codeHandler << If_ICmpGe(zero_label)
			case Not(v) =>
				generateExpressionCode(v)
				codeHandler << IfNe(zero_label)
			case And(lhs, rhs) =>
				generateExpressionCode(lhs)
				codeHandler << IfEq(zero_label)
				generateExpressionCode(rhs)
				codeHandler << IfEq(zero_label)
			case Or(lhs, rhs) =>
				val true_label = codeHandler.getFreshLabel("true_label")
				generateExpressionCode(lhs)
				codeHandler << IfNe(true_label) // we don't evaluate the second part (short-circuit)
				generateExpressionCode(rhs)
				codeHandler << IfEq(zero_label)
				codeHandler << Label(true_label)
			case _ => // shouldn't happen, but makes the compiler happy
		}

		codeHandler << Ldc(1)
		codeHandler << Goto(after_label)
		codeHandler << Label(zero_label)
		codeHandler << Ldc(0)
		codeHandler << Label(after_label)
	}

    def generateLoad(ident: Identifier) = ident.getType match {
        case t @ _ if isMember(ident) => codeHandler << ALoad(0) << GetField(currentClass.name, ident.value, generateTypeCodeSignature(t)) 
        case TInt | TBool => codeHandler << ILoad(getIndex(ident.value))
	    case TString | TArray => codeHandler << ALoad(getIndex(ident.value))
        case TObject(classSymbol) => codeHandler << ALoad(getIndex(ident.value))
		case _ => //shouldn't happen, but makes the compiler happy
    }

	/**
	 * Prerequisite : the value which is to be stored is already on top of the stack
	 */
	def generateStore(ident: Identifier, expr: ExprTree) = ident.getSymbol.asInstanceOf[VariableSymbol].getType match{
		case t @ _ if isMember(ident) =>
		    codeHandler << ALoad(0)
            generateExpressionCode(expr)
            codeHandler << PutField(currentClass.name,ident.value,generateTypeCodeSignature(t))
		case TInt | TBool =>
		    generateExpressionCode(expr);
            codeHandler << IStore(getIndex(ident.value))
		case TString | TArray | TObject(_) =>
		    generateExpressionCode(expr);
            codeHandler << AStore(getIndex(ident.value))
		case _ => 
	}
 
    def generateTypeCode(t: Type): String = t match {
        case TInt => "I"
		case TBool => "Z"
	    case TString => "Ljava/lang/String;"
	    case TArray => "[I"
        case TObject(value) => value.name // reverse-engineered value :)
        case _ => "error" // :-/
    }
    
    def generateTypeCodeSignature(t: Type): String = t match {
        case TObject(value) => "L"+value.name+";"
        case _ => generateTypeCode(t)
    }
}
